import colorsys
import os
import torch
import numpy as np
from scipy.interpolate import RegularGridInterpolator

LUT_FOLDER = os.path.join(os.path.dirname(__file__), "luts")


def __init__(self):
    pass

class OlmLUT:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "lut_file": (cls.get_lut_files(),),
                "strength": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01
                }),
                "gamma_correction": ("BOOLEAN", {
                    "default": False
                }),
                "debug_logging": ("BOOLEAN", {
                    "default": False
                }),
                "operation_mode": (["normal", "passthrough", "test_colors_horizontal", "test_colors_vertical", "test_gradient", "test_hsv", "mid_gray_box"],),
            }
        }

    RETURN_TYPES = ("IMAGE",)

    FUNCTION = "apply_lut"

    CATEGORY = "LUT"

    def load_cube_lut(self, filename, debug_logging):
        size = None
        lut_data = []

        lut_dir = os.path.join(os.path.dirname(__file__), "luts")
        path = os.path.join(lut_dir, filename)

        if debug_logging:
            print(f"Loading LUT from: {path}")

        try:
            with open(path, 'r') as f:
                for line in f:
                    line = line.strip()
                    if line == "" or line.startswith("#"):
                        continue
                    if "LUT_3D_SIZE" in line:
                        size = int(line.split()[-1])
                        if debug_logging:
                            print(f"LUT size: {size}")
                    elif all(c in "0123456789.+-eE " for c in line) and len(line.split()) == 3:
                        lut_data.append([float(v) for v in line.split()])
        except FileNotFoundError:
            raise ValueError(f"LUT file not found: {path}")

        if size is None:
            raise ValueError(f"No LUT_3D_SIZE found in LUT file: {path}")

        if len(lut_data) != size**3:
            raise ValueError(f"Expected {size**3} LUT entries, found {len(lut_data)} in file: {path}")

        lut_array = np.array(lut_data, dtype=np.float32)
        lut_array = lut_array.reshape(size, size, size, 3)
        lut_array = np.transpose(lut_array, (2, 1, 0, 3))

        if debug_logging:
            print(f"LUT loaded: shape {lut_array.shape}, range [{lut_array.min():.3f}, {lut_array.max():.3f}]")
            print(f"First few LUT entries:")
            for i in range(min(3, size)):
                print(f"  LUT[{i},0,0] = [{lut_array[i,0,0,0]:.3f}, {lut_array[i,0,0,1]:.3f}, {lut_array[i,0,0,2]:.3f}]")

        return lut_array

    @classmethod
    def get_lut_files(cls):
        folder = os.path.join(os.path.dirname(__file__), "luts")

        if not os.path.exists(folder):
            print(f"LUT folder does not exist: {folder}")
            return ["No LUTs folder found"]

        files = [f for f in os.listdir(folder) if f.lower().endswith(".cube")]
        return files if files else ["No LUTs found"]

    def srgb_to_linear(self, img):
        return np.where(img <= 0.04045,
                       img / 12.92,
                       np.power(np.clip((img + 0.055) / 1.055, 0, 1), 2.4))

    def linear_to_srgb(self, img):
        return np.where(img <= 0.0031308,
                       img * 12.92,
                       1.055 * np.power(np.clip(img, 0, 1), 1.0/2.4) - 0.055)

    def apply_lut(self, image, lut_file, strength, gamma_correction=False, operation_mode="normal", debug_logging=False):

        B, H, W, C = image.shape

        if C != 3:
            raise ValueError(f"Expected RGB image with 3 channels, got {C}")

        if debug_logging:
            print(f"\n=== LUT APPLICATION DEBUG ===")
            print(f"Input image shape: {image.shape}")
            print(f"Input range: [{image.min():.6f}, {image.max():.6f}]")
            print(f"Sample input pixel (0,0): R={image[0, 0, 0, 0]:.6f}, G={image[0, 0, 0, 1]:.6f}, B={image[0, 0, 0, 2]:.6f}")
            print(f"Gamma correction enabled: {gamma_correction}")
            print(f"Operation mode: {operation_mode}")

        test_image = [any]

        if operation_mode == "passthrough":
            if debug_logging:
                print("PASSTHROUGH MODE: Returning original image")
            return (image,)

        elif operation_mode == "test_colors_horizontal":
            if debug_logging:
                print("TEST COLORS MODE: Generating test pattern")

            test_image = torch.zeros_like(image)
            H, W = image.shape[1], image.shape[2]
            test_image[0, :H//3, :,       0] = 1.0
            test_image[0, H//3:2*H//3, :, 1] = 1.0
            test_image[0, 2*H//3:, :,     2] = 1.0

        elif operation_mode == "test_colors_vertical":
            if debug_logging:
                print("TEST COLORS VERTICAL MODE: Generating test pattern")

            test_image = torch.zeros_like(image)
            H, W = image.shape[1], image.shape[2]
            test_image[0, :, :W//3,       0] = 1.0
            test_image[0, :, W//3:2*W//3, 1] = 1.0
            test_image[0, :, 2*W//3:,     2] = 1.0

        elif operation_mode == "test_gradient":
            if debug_logging:
                print("TEST GRADIENT MODE: Generating test pattern")

            test_image = torch.zeros_like(image)
            H, W = image.shape[1], image.shape[2]
            for i in range(W):
                luma = i / (W - 1)
                test_image[0, :, i, 0] = luma
                test_image[0, :, i, 1] = luma
                test_image[0, :, i, 2] = luma

        elif operation_mode == "test_hsv":
            if debug_logging:
                print("TEST HSV MODE: Generating HSV hue sweep gradient")

            H, W = image.shape[1], image.shape[2]
            hues = np.linspace(0.0, 1.0, W, dtype=np.float32)
            rgb_row = np.array([colorsys.hsv_to_rgb(h, 1.0, 1.0) for h in hues], dtype=np.float32)  # shape: [W, 3]
            rgb_image = np.tile(rgb_row[np.newaxis, :, :], (H, 1, 1))  # shape: [H, W, 3]
            test_image = torch.from_numpy(rgb_image).unsqueeze(0)

        elif operation_mode == "mid_gray_box":
            if debug_logging:
                print("TEST GRAY BOX: Generating gray test box")

            lut = self.load_cube_lut(lut_file, debug_logging)
            lut_size = lut.shape[0]

            test_image = np.zeros((H, W, 3), dtype=np.float32)
            test_image[H//4:3*H//4, W//4:3*W//4, 0] = 0.5
            test_image[H//4:3*H//4, W//4:3*W//4, 1] = 0.5
            test_image[H//4:3*H//4, W//4:3*W//4, 2] = 0.5

            if debug_logging:
                mid_idx = lut_size // 2
                print(f"Mid-gray input [0.5, 0.5, 0.5] should map to LUT[{mid_idx},{mid_idx},{mid_idx}]:")
                print(f"  LUT output: [{lut[mid_idx,mid_idx,mid_idx,0]:.3f}, {lut[mid_idx,mid_idx,mid_idx,1]:.3f}, {lut[mid_idx,mid_idx,mid_idx,2]:.3f}]")

            test_image = torch.from_numpy(test_image).unsqueeze(0)

        output_images = []

        for b in range(B):

            if operation_mode in ["test_colors_horizontal", "test_colors_vertical", "test_gradient", "test_hsv", "mid_gray_box"]:
                img_np = test_image[b].cpu().numpy().astype(np.float32)
            else:
                img_np = image[b].cpu().numpy().astype(np.float32)

            if gamma_correction:
                img_for_lut = self.linear_to_srgb(img_np)
                if debug_logging:
                    print("Converting linear -> sRGB for LUT lookup")
                    print(f"Linear input sample: R={img_np[0,0,0]:.6f}, G={img_np[0,0,1]:.6f}, B={img_np[0,0,2]:.6f}")
                    print(f"sRGB for LUT sample: R={img_for_lut[0,0,0]:.6f}, G={img_for_lut[0,0,1]:.6f}, B={img_for_lut[0,0,2]:.6f}")
            else:
                if debug_logging:
                    print("No gamma correction - assuming input is already in LUT color space")
                img_for_lut = img_np.copy()

            lut = self.load_cube_lut(lut_file, debug_logging)
            lut_size = lut.shape[0]

            if debug_logging:
                print(f"\n--- LUT ANALYSIS ---")
                print(f"LUT shape: {lut.shape}")
                print(f"LUT range: [{lut.min():.6f}, {lut.max():.6f}]")

                print(f"LUT corners:")
                print(f"  Black [0,0,0] -> [{lut[0,0,0,0]:.3f}, {lut[0,0,0,1]:.3f}, {lut[0,0,0,2]:.3f}]")
                print(f"  White [{lut_size-1},{lut_size-1},{lut_size-1}] -> [{lut[-1,-1,-1,0]:.3f}, {lut[-1,-1,-1,1]:.3f}, {lut[-1,-1,-1,2]:.3f}]")
                mid = lut_size // 2
                print(f"  Mid-gray [{mid},{mid},{mid}] -> [{lut[mid,mid,mid,0]:.3f}, {lut[mid,mid,mid,1]:.3f}, {lut[mid,mid,mid,2]:.3f}]")

            coords = [np.linspace(0, 1, lut_size) for _ in range(3)]

            interpolators = []
            for channel in range(3):
                interpolators.append(
                    RegularGridInterpolator(
                        coords,
                        lut[:, :, :, channel],
                        method='linear',
                        bounds_error=False,
                        fill_value=0.0
                    )
                )

            img_flat = img_for_lut.reshape(-1, 3)
            img_flat = np.clip(img_flat, 0, 1)

            mapped_flat = np.zeros_like(img_flat, dtype=np.float32)

            chunk_size = 10000
            for i in range(0, len(img_flat), chunk_size):
                end_idx = min(i + chunk_size, len(img_flat))
                chunk = img_flat[i:end_idx]

                for channel in range(3):
                    mapped_flat[i:end_idx, channel] = interpolators[channel](chunk)

            mapped_img = mapped_flat.reshape(H, W, 3)

            if gamma_correction:
                mapped_img_linear = self.srgb_to_linear(mapped_img)
                if debug_logging:
                    print("Converting sRGB LUT result -> linear")
                    print(f"sRGB LUT result sample: R={mapped_img[0,0,0]:.6f}, G={mapped_img[0,0,1]:.6f}, B={mapped_img[0,0,2]:.6f}")
                    print(f"Linear result sample: R={mapped_img_linear[0,0,0]:.6f}, G={mapped_img_linear[0,0,1]:.6f}, B={mapped_img_linear[0,0,2]:.6f}")
                mapped_img = mapped_img_linear

            blended = (1 - strength) * img_np + strength * mapped_img
            blended = np.clip(blended, 0, 1).astype(np.float32)

            if debug_logging:
                print(f"Original vs LUT comparison for sample pixel (center of the image):")
                print(f"  Original: R={img_np[H//2,W//2,0]:.6f}, G={img_np[H//2,W//2,1]:.6f}, B={img_np[H//2,W//2,2]:.6f}")
                print(f"  LUT result: R={mapped_img[H//2,W//2,0]:.6f}, G={mapped_img[H//2,W//2,1]:.6f}, B={mapped_img[H//2,W//2,2]:.6f}")
                print(f"  Final blend: R={blended[H//2,W//2,0]:.6f}, G={blended[H//2,W//2,1]:.6f}, B={blended[H//2,W//2,2]:.6f}")

            output_images.append(blended)

        output_tensor = torch.from_numpy(np.stack(output_images, axis=0))

        if debug_logging:
            print(f"\n=== FINAL OUTPUT ===")
            print(f"Output shape: {output_tensor.shape}")
            print(f"Output range: [{output_tensor.min():.6f}, {output_tensor.max():.6f}]")
            print("=== END DEBUG ===\n")

        return (output_tensor,)


NODE_CLASS_MAPPINGS = {
    "OlmLUT": OlmLUT,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "OlmLUT": "Olm LUT",
}